package org.bdware.irp.irplib.util;

import org.bdware.irp.irplib.exception.IrpMessageCredentialException;

import com.nimbusds.jose.*;
import com.nimbusds.jose.crypto.*;
import com.nimbusds.jose.jwk.Curve;
import com.nimbusds.jose.jwk.JWK;
import com.nimbusds.jose.jwk.KeyType;
import com.nimbusds.jose.jwk.KeyUse;
import com.nimbusds.jose.jwk.gen.ECKeyGenerator;

import org.apache.log4j.Logger;

import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.text.ParseException;

public class GlobalUtils {
    private static Logger logger = Logger.getLogger(GlobalUtils.class);

    //encode string to bytes by specified coding type
    public static final byte[] encodeString(String str) {
        try {
            return str.getBytes(IrpCommon.MESSAGE_ENCODING_TYPE);
        } catch (Exception e) {
            logger.error(str + " encode to bytes failed!");
        }
        return null;
    }

    //encode bytes to String by specified coding type
    public static final String decodeString(byte[] bytes) {
        if (bytes == null | bytes.length == 0 )
            return "";
        try {
            return new String(bytes, IrpCommon.MESSAGE_ENCODING_TYPE);
        } catch (Exception e) {
            logger.error(new String(bytes) + " decode to string failed!");
        }
        return null;
    }

    public static final byte[] getDigestAlgFromSignature(String sigAlgorithm) throws IrpMessageCredentialException {
        if (sigAlgorithm.startsWith("SHA1"))
            return IrpCommon.CREDENTIAL_DIGEST_ALG_SHA1;
        else if (sigAlgorithm.startsWith("MD5"))
            return IrpCommon.CREDENTIAL_DIGEST_ALG_MD5;
        else
            throw new IrpMessageCredentialException("Unsupported digest algorithm for signature: " + sigAlgorithm);
    }

    public static final String signByteArrayByJWK(byte[] data, JWK jwk) throws JOSEException{
        JWSSigner jwsSigner;
        JWSObject jwsObject;

        if(jwk.getKeyType() == KeyType.RSA){
            jwsSigner = new RSASSASigner(jwk.toRSAKey());
            jwsObject = new JWSObject(
                new JWSHeader.Builder(JWSAlgorithm.RS256).keyID(jwk.getKeyID()).build(),
                new Payload(data));
        }else if(jwk.getKeyType() == KeyType.EC){
            jwsSigner = new ECDSASigner(jwk.toECKey());
            jwsObject = new JWSObject(
                new JWSHeader.Builder(JWSAlgorithm.ES256).keyID(jwk.getKeyID()).build(),
                new Payload(data));
        }else{
            logger.error("unsupported jwk Algorithm");
            return null;
        }
        //Compute the EC signature
        jwsObject.sign(jwsSigner);

        return jwsObject.serialize(true);
    }

    public static final boolean verifySigByJWK(byte[] messageBody, String signature, String pkInfo){
        if (messageBody == null || signature == null || pkInfo == null){
            logger.error("SignatureInfo missing!");
            return false;
        }
        try{
            JWSObject verify = JWSObject.parse(signature, new Payload(messageBody));
            //parse public key string
            JWK pkToVerify = JWK.parse(pkInfo);
            JWSVerifier verifier;
            if(pkToVerify.getKeyType() == KeyType.RSA){
                verifier = new RSASSAVerifier(pkToVerify.toRSAKey().toRSAPublicKey());
            }else if(pkToVerify.getKeyType() == KeyType.EC){
                verifier = new ECDSAVerifier(pkToVerify.toECKey().toECPublicKey());
            }else{
                logger.error("unsupported Algorithm");
                return false;
            }
            if(verify.verify(verifier)){
                return true;
            }else{
                logger.error("verify the message failed!");
                return false;
            }
        }catch(ParseException e){
            logger.error("parse the pk error");
            return false;
        }catch(JOSEException e) {
            e.printStackTrace();
            return false;
        }
    }

    public static JWK loadKeysFromJWKFile(String jwkFilePath) throws IOException, ParseException {
        FileInputStream fi = new FileInputStream(jwkFilePath);
        byte[] jwkB = new byte[fi.available()];
        fi.read(jwkB);
        String jwkS = new String(jwkB);
        logger.info(jwkS);
        return JWK.parse(jwkS);
    }
}
